#include <vector>
#include <unordered_map>
#include <algorithm>

#include "constants.h"
#include "classes.h"
#include "tools.h"
#include "optimizer_tools.h"
#include "optimizer_radical.h"
using namespace std;

OptRadical::OptRadical(AssignInfo &assign_info)
{
    this->assign_info = &assign_info;
    vector<BwType> _tmp(N, LOW);
    this->edge_site_bw_type_vec.reserve(T);
    this->center_site_bw_type_vec.reserve(T);
    for (int t = 0; t < T; ++t)
    {
        this->edge_site_bw_type_vec.push_back(_tmp);
        this->center_site_bw_type_vec.push_back(LOW);
    }
}

void OptRadical::opt(double edge_low_quantile, double edge_hight_quantile,
                     double center_low_quantile, double center_hight_quantile,
                     int skip_diff_sum)
{
    this->skip_diff_sum = skip_diff_sum;

    this->edge_low_index = ceil(edge_low_quantile * T) - 1;
    this->edge_hight_index = ceil(edge_hight_quantile * T) - 1;
    this->center_low_index = ceil(center_low_quantile * T) - 1;
    this->center_hight_index = ceil(center_hight_quantile * T) - 1;
    this->edge_site_score = 0.0;
    this->center_site_score = 0.0;
    this->success_count = 0;

    this->refresh_state();
    this->total_score = this->edge_site_score + this->center_site_score;

    cout << "total_score:" << (long long)this->total_score
         << " edge_site_score:" << (long long)this->edge_site_score
         << " center_site_score:" << (long long)this->center_site_score
         << endl;

    for (int t = 0; t < T; t++)
    {
        this->opt_t(t);
    }
    cout << "success transfer radical count:" << this->success_count << endl;
}

void OptRadical::refresh_state()
{
    for (int site_id = 0; site_id < N; ++site_id)
    {
        this->refresh_edge_site_state(site_id);
    }
    this->refresh_center_site_state();
}

void OptRadical::refresh_edge_site_state(int site_id)
{
    vector<TAssign *> &assign_vec = this->assign_info->assign_vec;
    int bw_allo_list[T];
    for (int t = 0; t < T; ++t)
    {
        bw_allo_list[t] = assign_vec[t]->site_total_bw[site_id];
    }
    vector<int> bw_sorted_index = arg_sort(bw_allo_list, T, true);

    int target_bw = bw_allo_list[bw_sorted_index[config_info.quantile_index]];
    this->edge_site_score += get_edge_site_score(target_bw, site_info.site_bw_upper[site_id]);

    this->edge_site_low_bw[site_id] = max(bw_allo_list[bw_sorted_index[this->edge_low_index]], config_info.base_cost);
    this->edge_site_high_bw[site_id] = max(bw_allo_list[bw_sorted_index[this->edge_hight_index]], config_info.base_cost);

    for (int i = 0; i < this->edge_low_index; ++i)
    {
        this->edge_site_bw_type_vec[bw_sorted_index[i]][site_id] = LOW;
    }

    for (int i = this->edge_low_index; i <= this->edge_hight_index; ++i)
    {
        this->edge_site_bw_type_vec[bw_sorted_index[i]][site_id] = CRUCIAL;
    }

    for (int i = this->edge_hight_index + 1; i < T; ++i)
    {
        this->edge_site_bw_type_vec[bw_sorted_index[i]][site_id] = BUFFER;
    }

    int _T = T - 1;
    for (int t = 0; t < _T; ++t)
    {
        if ((this->edge_site_bw_type_vec[t][site_id] == BUFFER) && (this->edge_site_bw_type_vec[t + 1][site_id] == BUFFER))
        {
            this->edge_site_bw_type_vec[t][site_id] = RADICAL;
        }
    }
    if (this->edge_site_bw_type_vec[_T][site_id] == BUFFER)
    {
        this->edge_site_bw_type_vec[_T][site_id] = RADICAL;
    }
}

void OptRadical::refresh_center_site_state()
{
    vector<TAssign *> &assign_vec = this->assign_info->assign_vec;
    long long bw_allo_list[T];
    for (int t = 0; t < T; ++t)
    {
        bw_allo_list[t] = assign_vec[t]->center_site_bw;
    }
    vector<int> bw_sorted_index = arg_sort(bw_allo_list, T, true);

    long long target_bw = bw_allo_list[bw_sorted_index[config_info.quantile_index]];
    this->center_site_score = get_center_site_score(target_bw);

    this->center_site_low_bw = bw_allo_list[bw_sorted_index[this->center_low_index]];
    this->center_site_high_bw = bw_allo_list[bw_sorted_index[this->center_hight_index]];

    for (int i = 0; i < this->center_low_index; ++i)
    {
        this->center_site_bw_type_vec[bw_sorted_index[i]] = LOW;
    }

    for (int i = this->center_low_index; i <= this->center_hight_index; ++i)
    {
        this->center_site_bw_type_vec[bw_sorted_index[i]] = CRUCIAL;
    }

    for (int i = this->center_hight_index + 1; i < T; ++i)
    {
        this->center_site_bw_type_vec[bw_sorted_index[i]] = RADICAL;
    }
}

void OptRadical::opt_t(int t)
{
    TAssign *tassign = this->assign_info->assign_vec[t];

    tassign->refresh_inherited_bw();
    tassign->refresh_total_bw();
    process_overflow_t(tassign);

    vector<BwType> &edge_site_bw_types = this->edge_site_bw_type_vec[t];

    for (int site_id = 0; site_id < N; site_id++)
    {
        BwType &site_bw_type = edge_site_bw_types[site_id];
        if ((site_bw_type == BUFFER) || (site_bw_type == RADICAL))
        {
            this->opt_t_big_site(t, tassign, site_id);
        }
    }
}

void OptRadical::opt_t_big_site(int t, TAssign *tassign, int site_id)
{
    int target = min(this->edge_site_low_bw[site_id] - 1, site_info.site_bw_upper[site_id]);
    this->success_count += this->opt_t_big_site_try_decrease(t, tassign, site_id, target, false);
}

bool OptRadical::opt_t_big_site_try_decrease(int t, TAssign *tassign, int site_id, int target, bool is_recursion_call)
{
    if (target <= 0)
    {
        return false;
    }

    vector<int> safe_bw_uppers = this->get_safe_bw_uppers(t, tassign);
    safe_bw_uppers[site_id] = target;

    int diff_sum = 0;
    for (int site_id = 0; site_id < N; site_id++)
    {
        diff_sum += (safe_bw_uppers[site_id] - tassign->site_total_bw[site_id]);
    }
    if (diff_sum < this->skip_diff_sum)
    {
        return false;
    }

    if (tassign->site_inherited_bw[site_id] > target)
    {
        return false;
        // if (is_recursion_call)
        // {
        //     return false;
        // }
        // else
        // {
        //     if (this->opt_t_big_site_try_decrease(t - 1, tassign->pre_tassign, site_id, this->edge_site_low_bw[site_id] / INHERIT_RATE, true))
        //     {
        //         tassign->refresh_inherited_bw();
        //         tassign->refresh_total_bw();
        //         tassign->snapshot();
        //         if (this->opt_t_big_site_process_overflow_t(tassign, safe_bw_uppers))
        //         {
        //             tassign->pre_tassign->cancel_snapshot();
        //             tassign->cancel_snapshot();
        //             this->edge_site_bw_type_vec[t][site_id] = CRUCIAL;
        //             return true;
        //         }
        //         else
        //         {
        //             tassign->restore();
        //             tassign->pre_tassign->restore();

        //             tassign->refresh_inherited_bw();
        //             tassign->refresh_total_bw();
        //             return false;
        //         }
        //     }
        //     else
        //     {
        //         tassign->pre_tassign->restore();

        //         tassign->refresh_inherited_bw();
        //         tassign->refresh_total_bw();
        //         return false;
        //     }
        // }
    }
    else
    {
        tassign->snapshot();
        if (this->opt_t_big_site_process_overflow_t(tassign, safe_bw_uppers))
        {
            if (!is_recursion_call)
            {
                tassign->cancel_snapshot();
                this->edge_site_bw_type_vec[t][site_id] = CRUCIAL;
            }
            return true;
        }
        else
        {
            if (!is_recursion_call)
            {
                tassign->restore();
            }
            return false;
        }
    }
    return false;
}

bool OptRadical::opt_t_big_site_process_overflow_t(TAssign *tassign, vector<int> &bw_uppers)
{
    for (int site_id = 0; site_id < N; site_id++)
    {
        if (tassign->site_total_bw[site_id] > bw_uppers[site_id])
        {
            SiteAlloSortIter site_allo_sort_iter(tassign->site_stream_to_allo_info[site_id], true);
            AlloNode *min_node;
            site_allo_sort_iter.next(min_node);
            // while ((min_node) && (tassign->site_total_bw[site_id] > bw_uppers[site_id]))
            while (min_node)
            {
                const int &client_id = min_node->client_id;
                vector<int> &selectable_site_ids = client_info.client_to_sites[client_id];

                int transfer_site_id = site_id;
                for (auto &adj_site_id : selectable_site_ids)
                {
                    if ((tassign->site_total_bw[adj_site_id] + min_node->bandwidth) <= bw_uppers[adj_site_id])
                    {
                        transfer_site_id = adj_site_id;
                        break;
                    }
                }

                if (transfer_site_id != site_id)
                {
                    tassign->transfer(min_node->stream_id, client_id, transfer_site_id);
                }

                site_allo_sort_iter.next(min_node);
            }
        }
    }

    for (int site_id = 0; site_id < N; site_id++)
    {
        if (tassign->site_total_bw[site_id] > bw_uppers[site_id])
        {
            return false;
        }
    }
    return true;
};

vector<int> OptRadical::get_safe_bw_uppers(int t, TAssign *tassign)
{
    vector<int> bw_upper(N, 0);
    vector<BwType> &edge_site_bw_types = this->edge_site_bw_type_vec[t];
    for (int site_id = 0; site_id < N; site_id++)
    {
        const int &low_bw = this->edge_site_low_bw[site_id];
        int &_upper = bw_upper[site_id];
        const BwType &bw_type = edge_site_bw_types[site_id];
        switch (bw_type)
        {
        case LOW:
            _upper = low_bw;
            break;
        case CRUCIAL:
            _upper = max(tassign->site_total_bw[site_id], config_info.base_cost);
            break;
        case BUFFER: // BUFFER一定不会出现在最后一个时刻
        {
            _upper = (low_bw - tassign->next_tassign->site_allocated_bw[site_id]) / INHERIT_RATE;
            _upper = max(_upper, tassign->site_total_bw[site_id]);
            break;
        }
        case RADICAL:
        {
            if (tassign->next_tassign)
            {
                _upper = (site_info.site_bw_upper[site_id] - tassign->next_tassign->site_allocated_bw[site_id]) / INHERIT_RATE;
                _upper = max(_upper, tassign->site_total_bw[site_id]);
            }
            else
            {
                _upper = site_info.site_bw_upper[site_id];
            }
            break;
        }
        }
        _upper = min(_upper, site_info.site_bw_upper[site_id]);
    }
    return bw_upper;
}
